import numpy as np
import inspect
import functools


def store_args(method):
    """Stores provided method args as instance attributes."""
    argspec = inspect.getfullargspec(method)
    defaults = {}
    if argspec.defaults is not None:
        defaults = dict(
            zip(argspec.args[-len(argspec.defaults):], argspec.defaults))
    if argspec.kwonlydefaults is not None:
        defaults.update(argspec.kwonlydefaults)
    arg_names = argspec.args[1:]

    @functools.wraps(method)
    def wrapper(*positional_args, **keyword_args):
        self = positional_args[0]
        # Get default arg values
        args = defaults.copy()
        # Add provided arg values
        for name, value in zip(arg_names, positional_args[1:]):
            args[name] = value
        args.update(keyword_args)
        self.__dict__.update(args)
        return method(*positional_args, **keyword_args)

    return wrapper


# 在make_encirclement_env函数中添加种子参数
def make_encirclement_env(args):
    """创建合围环境"""
    from environment.encirclement_env import EncirclementEnv
    env = EncirclementEnv(args.num_uavs, args.grid_size, 
                         args.num_position_action, args.num_energy_action, 
                         seed=args.seed)  # 传递种子参数
    return env, args



def make_uav_env(args):
    from environment.environment import MultiUAVEnv
    env = MultiUAVEnv(args.num_uavs, args.grid_size, args.num_users,args.num_position_action,args.num_energy_action)
    return env, args